import sys
import torch

def check_versions():
    # 获取 Python 版本
    python_version = sys.version.split(" ")[0]
    print(f"Python 版本: {python_version}")

    # 检查 CUDA 是否可用
    if torch.cuda.is_available():
        # 获取 CUDA 版本
        cuda_version = torch.version.cuda
        print(f"CUDA 版本: {cuda_version}")

        # 获取 cuDNN 版本
        cudnn_version = torch.backends.cudnn.version()
        print(f"cuDNN 版本: {cudnn_version}")
    else:
        print("CUDA 不可用，无法获取 CUDA 和 cuDNN 版本。")

if __name__ == "__main__":
    check_versions()